#ifndef RSAUTIL_H
#define RSAUTIL_H

#include <memory>
#include <openssl/err.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <stdexcept>
#include <string>
#include <vector>

class RSAUtil
{
	public:
	// 使用公钥加密
	static std::string encrypt(const std::string &original_text, const std::string &public_key_file_path)
	{
		EVP_PKEY *public_key = loadPublicKey(public_key_file_path);
		auto ctx = std::unique_ptr<EVP_PKEY_CTX, decltype(&EVP_PKEY_CTX_free)>(EVP_PKEY_CTX_new(public_key, nullptr),
																			   EVP_PKEY_CTX_free);
		if (!ctx) throw std::runtime_error("Failed to create encryption context");

		if (EVP_PKEY_encrypt_init(ctx.get()) <= 0) throw std::runtime_error("Encrypt init failed");

		if (EVP_PKEY_CTX_set_rsa_padding(ctx.get(), RSA_PKCS1_OAEP_PADDING) <= 0)
			throw std::runtime_error("Failed to set OAEP padding");

		size_t key_len = EVP_PKEY_get_size(public_key);
		const size_t HASH_LEN = 32; // SHA-256
		const size_t MAX_BLOCK_SIZE = key_len - 2 * HASH_LEN - 2;

		std::string encrypted_data;

		for (size_t i = 0; i < original_text.size(); i += MAX_BLOCK_SIZE) {
			size_t block_size = std::min(MAX_BLOCK_SIZE, original_text.size() - i);
			std::string block = original_text.substr(i, block_size);

			size_t outlen;
			if (EVP_PKEY_encrypt(ctx.get(), nullptr, &outlen, reinterpret_cast<const unsigned char *>(block.data()),
								 block.size()) <= 0) {
				throw std::runtime_error("Encryption buffer size calculation failed");
			}

			std::vector<unsigned char> encrypted_block(outlen);
			if (EVP_PKEY_encrypt(ctx.get(), encrypted_block.data(), &outlen,
								 reinterpret_cast<const unsigned char *>(block.data()), block.size()) <= 0) {
				throw std::runtime_error("Encryption failed for block");
			}

			// ✅ 只追加有效长度
			encrypted_data.append(reinterpret_cast<char *>(encrypted_block.data()), outlen);
		}

		EVP_PKEY_free(public_key);
		return encrypted_data;
	}

	// 使用私钥解密
	static std::string decrypt(const std::string &encrypted_text, const std::string &private_key_file_path)
	{
		EVP_PKEY *private_key = loadPrivateKey(private_key_file_path);
		auto ctx = std::unique_ptr<EVP_PKEY_CTX, decltype(&EVP_PKEY_CTX_free)>(EVP_PKEY_CTX_new(private_key, nullptr),
																			   EVP_PKEY_CTX_free);
		if (!ctx) throw std::runtime_error("Failed to create decryption context");

		if (EVP_PKEY_decrypt_init(ctx.get()) <= 0) throw std::runtime_error("Decrypt init failed");

		if (EVP_PKEY_CTX_set_rsa_padding(ctx.get(), RSA_PKCS1_OAEP_PADDING) <= 0)
			throw std::runtime_error("Failed to set OAEP padding");

		size_t key_len = EVP_PKEY_get_size(private_key);

		std::string decrypted_data;

		for (size_t i = 0; i < encrypted_text.size(); i += key_len) {
			size_t block_size = std::min(key_len, encrypted_text.size() - i);
			std::string block = encrypted_text.substr(i, block_size);

			size_t outlen;
			if (EVP_PKEY_decrypt(ctx.get(), nullptr, &outlen, reinterpret_cast<const unsigned char *>(block.data()),
								 block.size()) <= 0) {
				throw std::runtime_error("Decryption buffer size calculation failed");
			}

			std::vector<unsigned char> decrypted_block(outlen);
			if (EVP_PKEY_decrypt(ctx.get(), decrypted_block.data(), &outlen,
								 reinterpret_cast<const unsigned char *>(block.data()), block.size()) <= 0) {
				throw std::runtime_error("Decryption failed for block");
			}

			// ✅ 只追加有效长度
			decrypted_data.append(reinterpret_cast<char *>(decrypted_block.data()), outlen);
		}

		EVP_PKEY_free(private_key);
		return decrypted_data;
	}

	private:
	static EVP_PKEY *loadPublicKey(const std::string &file_path)
	{
		FILE *file = fopen(file_path.c_str(), "rb");
		if (!file) throw std::runtime_error("Cannot open public key file");

		EVP_PKEY *key = PEM_read_PUBKEY(file, nullptr, nullptr, nullptr);
		fclose(file);

		if (!key) throw std::runtime_error("Invalid public key format");
		return key;
	}

	static EVP_PKEY *loadPrivateKey(const std::string &file_path)
	{
		FILE *file = fopen(file_path.c_str(), "rb");
		if (!file) throw std::runtime_error("Cannot open private key file");

		EVP_PKEY *key = PEM_read_PrivateKey(file, nullptr, nullptr, nullptr);
		fclose(file);

		if (!key) throw std::runtime_error("Invalid private key format");
		return key;
	}
};

#endif
